package com.walter.gateway.filters;

import com.walter.gateway.utils.IpAddrUtil;
import io.netty.buffer.ByteBufAllocator;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.util.StopWatch;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

/**
 * @author Walter
 * @date 2025/2/5 下午4:15
 **/
@Slf4j
@Component
public class RequestHandledFilter implements GlobalFilter, Ordered {

    private static final String REQUEST_BODY_START_TIME = "requestBodyStartTime";

    private final DataBufferFactory dataBufferFactory = new NettyDataBufferFactory(ByteBufAllocator.DEFAULT);

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        StopWatch stopWatch = new StopWatch();
        stopWatch.start();
        ServerHttpRequest request = exchange.getRequest();
        HttpMethod method = request.getMethod();
        String path = request.getPath().value();
        String ip = IpAddrUtil.getClientIpAddress(request);
        Map<String, String> queryParams = this.getQueryParameters(request);

        // 打印请求 Query 参数
        if (!queryParams.isEmpty()) {
            log.info("Request path:{} Query Params: {}", path, queryParams);
        }
        Mono<Void> filterResult;
        // 获取 Request Body (只针对 POST/PUT/PATCH 等有 body 的请求)
        if (method == HttpMethod.POST || method == HttpMethod.PUT || method == HttpMethod.PATCH) {
            filterResult = DataBufferUtils.join(exchange.getRequest().getBody())
                    .flatMap(dataBuffer -> {
                        byte[] content = new byte[dataBuffer.readableByteCount()];
                        dataBuffer.read(content);
                        DataBufferUtils.release(dataBuffer);
                        String requestBody = new String(content, StandardCharsets.UTF_8);
                        // 打印 Request Body
                        log.info("Request path:{} Body: {}", path, requestBody);
                        // 重新构建 Request，将 body 塞回去，否则下游服务会收不到 body
                        ServerHttpRequest modifiedRequest = new ServerHttpRequestDecorator(exchange.getRequest()) {
                            @Override
                            public Flux<DataBuffer> getBody() {
                                DataBuffer dataBuffer = dataBufferFactory.allocateBuffer();
                                dataBuffer.write(content);
                                return Flux.just(dataBuffer);
//                                return Flux.just(exchange.getRequest().bufferFactory().wrap(content));
                            }
                        };
                        exchange.getAttributes().put(REQUEST_BODY_START_TIME, System.currentTimeMillis());
                        return chain.filter(exchange.mutate().request(modifiedRequest).build());

                    });
        } else {
            exchange.getAttributes().put(REQUEST_BODY_START_TIME, System.currentTimeMillis());
            filterResult = chain.filter(exchange);
        }
        return filterResult.then(Mono.fromRunnable(() -> {
            stopWatch.stop();
            ServerHttpResponse response = exchange.getResponse();
            long startTime = exchange.getAttributeOrDefault(REQUEST_BODY_START_TIME, System.currentTimeMillis());
            long endTime = System.currentTimeMillis();
            long duration = endTime - startTime;
            // 获取响应状态码
            int statusCode = response.getStatusCode() != null ? response.getStatusCode().value() : -1;

            // 获取 Response Body  (注意： 生产环境需要考虑性能和 body 大小限制)
            ServerHttpResponseDecorator responseDecorator = new ServerHttpResponseDecorator(response) {
                @Override
                public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                    if (body instanceof Flux) {
                        Flux<? extends DataBuffer> fluxBody = (Flux<? extends DataBuffer>) body;
                        return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                            // 合并多个 DataBuffer
                            DataBufferFactory dataBufferFactory = response.bufferFactory();
                            DataBuffer joinedBuffer = dataBufferFactory.allocateBuffer();
                            dataBuffers.forEach(joinedBuffer::write);

                            byte[] content = new byte[joinedBuffer.readableByteCount()];
                            joinedBuffer.read(content);
                            DataBufferUtils.release(joinedBuffer);
                            String responseBody = new String(content, StandardCharsets.UTF_8);
                            // 打印 Response Body
                            log.info("Response Body: {}", responseBody);
                            // 重新构建 DataBuffer 返回，重要！ 否则响应内容就丢失了
                            return dataBufferFactory.wrap(content);
                        }));
                    }
                    // if body is not a flux. never got there.
                    return super.writeWith(body);
                }
            };
            log.info("clientIp: {}, Method: {}, Path: {}, Status Code: {}, Duration: {}ms", ip, method, path, statusCode, duration);
        }));
    }

    @Override
    public int getOrder() {
        return Ordered.LOWEST_PRECEDENCE;
    }

    /**
     * 获取请求的 Query 参数
     * @param request
     * @return
     */
    private Map<String, String> getQueryParameters(ServerHttpRequest request) {
        Map<String, String> queryParamsMap = new HashMap<>();
        request.getQueryParams().forEach((name, values) -> {
            if (!values.isEmpty()) {
                queryParamsMap.put(name, String.join(",", values)); // 如果有多个值，用逗号分隔
            }
        });
        return queryParamsMap;
    }

}
